package com.example.module.ai.spark.websocket;

import com.example.framework.common.exception.ServiceException;
import com.example.framework.common.util.json.JsonUtils;
import com.example.module.ai.spark.dto.GptRespDto;
import lombok.Getter;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import okhttp3.Response;
import okhttp3.WebSocket;
import okhttp3.WebSocketListener;
import org.jetbrains.annotations.NotNull;

import java.util.List;
import java.util.Optional;
import java.util.concurrent.locks.Condition;
import java.util.concurrent.locks.Lock;
import java.util.concurrent.locks.ReentrantLock;

@RequiredArgsConstructor
@Slf4j
@Getter
public class SparkWebSocketMessageListener extends WebSocketListener {

    private static final int MAX_HISTORY_COUNT = 12000;

    // 大模型的答案汇总
    private final StringBuilder totalAnswer = new StringBuilder();

    // 大模型调用错误标识
    private boolean errFlag = false;
    // 大模型调用错误信息
    private String errMsg = null;

    // 发送问题和等待回答的协调机制，lock和condition方式
    private final Lock lock = new ReentrantLock();
    private final Condition condition = lock.newCondition();

    // 发送问题和等待回答的协调机制，wait/notify方式
//    private static final Object totalFlag = new Object(); // 消息是否已回答完

    private final List<GptRespDto.RoleContent> historyList;

    @Override
    public void onMessage(@NotNull WebSocket webSocket, @NotNull String text) {
        log.debug("星火返回消息：{}", text);
        GptRespDto.JsonParse myJsonParse = JsonUtils.parseObject(text, GptRespDto.JsonParse.class);
        Optional.ofNullable(myJsonParse).orElseThrow(() -> ServiceException.of("星火返回消息解析失败"));
        if (myJsonParse.getHeader().getCode() != GptRespDto.Header.SUCCESS_CODE) {
            log.error("发生错误，errCode: {}, sid：{}", myJsonParse.getHeader().getCode(), myJsonParse.getHeader().getSid());
            errFlag = true;
            errMsg = text;
            closeWebSocketAndSignalCondition(webSocket);
            return;
        }
        List<GptRespDto.Text> textList = myJsonParse.getPayload().getChoices().getText();
        for (GptRespDto.Text temp : textList) {
            totalAnswer.append(temp.getContent());
        }
        if (myJsonParse.getHeader().getStatus() == GptRespDto.Header.LAST_SESSION_STATUS) {
            addHistory();
            // 可以关闭连接，释放资源
            closeWebSocketAndSignalCondition(webSocket);
        }
    }

    private void addHistory() {
        GptRespDto.RoleContent roleContent = new GptRespDto.RoleContent()
                .setRole("assistant")
                .setContent(totalAnswer.toString());
        historyList.add(roleContent);

        // 由于历史记录最大上线1.2W左右，需要判断是否能加入历史
        int historyLength = historyList.stream()
                .mapToInt(item -> item.getContent().length())
                .sum();
        while (historyLength > MAX_HISTORY_COUNT && !historyList.isEmpty()) {
            GptRespDto.RoleContent removedItem = historyList.removeFirst();
            historyLength -= removedItem.getContent().length();
        }
    }

    private void closeWebSocketAndSignalCondition(WebSocket webSocket) {
        lock.lock();
        condition.signal();
        lock.unlock();
        // wait/notify方式
//            synchronized (totalFlag) {
//                totalFlag.notify();
//            }
        webSocket.close(1000, "");
    }

    @Override
    public void onFailure(@NotNull WebSocket webSocket, @NotNull Throwable t, Response response) {
        super.onFailure(webSocket, t, response);
        if (null != response) {
            int code = response.code();
            log.error("onFailure, code: {}, body: {}", code, response.body());
            if (101 != code) {
                log.error("connection failed");
                System.exit(0);
            }
        }
    }
}
